#! /usr/bin/env python
# coding=utf-8

# 向目标 mRNA 和pep 文件 添加 参考物种的mRNA和pep
# 需要参考物种的 mRNA文件



import string
from Bio import SeqIO
from Bio import Seq

import sys
import argparse



dis = '''
将参考物种的mRNA 和pep序列添加回去
tblastn_fetch_gene01_adding_reference_spes.py -r ref_spe_table -m mRNA_raw_ali.fa  -c filter_mRNA.fa -p filter_add 
'''


parser = argparse.ArgumentParser(description=dis)



parser.add_argument('-r',
                help='必须给定，参考物种的列表')

parser.add_argument('-m',
                help='必须给定，包含参考物种的原始mRNA文件')

parser.add_argument('-p',
                help='必须给定，结果前缀')

parser.add_argument('-c',
                help='必须给定，需要添加的过滤完科的mRNA文件 ')




# 获取参数
args = parser.parse_args()



if not args.p or not args.m or not args.c or not args.r:
    parser.print_help()
    sys.exit()



ref_spe_infile = args.r #r'ref_spe_table'

ref_mRNA_file = args.m # r'mRNA_raw_ali.fa'

current_mRNA_file = args.c #'filter_mRNA.fa'

out_mRNA_file = args.p+'_mRNA.fa'

out_pep_file = args.p+'_pep.fa'



out_mRNA_file = open(out_mRNA_file,'w')
out_pep_file = open(out_pep_file,'w')


abb_dic = {}
genus_order = []

with open(ref_spe_infile) as fila:
	for i in fila:
		k = i.strip().split('\t')
		if len(k)>2:
			abb = k[3]
			family = k[1]
			genus =k[0].split(' ')[0]
			if genus not in genus_order:
				genus_order.append(genus)
			spe = k[0]
			order = k[2]
			new_abb = string.capwords(order[:2])+string.capwords(family[:3])+abb
			abb_dic[abb] = {'genus':genus,'order':order,'family':family,'spe':spe,'new_abb':new_abb}

#target_seq_pep = {}



for i in SeqIO.parse(ref_mRNA_file,'fasta'):
	if str(i.name) in abb_dic:
		mRNA = str(i.seq).replace('-','')
		pep = str(Seq.Seq(mRNA).translate())

		if pep.strip()!='' and mRNA.strip()!='':
		#target_seq_pep[] = (mRNA,pep)
			out_pep_file.write('>'+abb_dic[i.name]['new_abb']+'\n'+pep+'\n')
			out_mRNA_file.write('>'+abb_dic[i.name]['new_abb']+'\n'+mRNA+'\n')

# 将目标序列添加到 新序列中

#current_mRNA_dic = {}
for i in SeqIO.parse(current_mRNA_file,'fasta'):
	mRNA = str(i.seq).replace('-','')
	pep = str(Seq.Seq(mRNA).translate())
	if pep.strip()!='' and mRNA.strip()!='':

		out_pep_file.write('>'+str(i.name)+'\n'+pep+'\n')
		out_mRNA_file.write('>'+str(i.name)+'\n'+mRNA+'\n')





out_mRNA_file.close()
out_pep_file.close()
